// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/text/writer/buffered_writer_sink.h>

namespace kumo::pollux::text {
    BufferedWriterSink::BufferedWriterSink(
        std::unique_ptr<dwio::common::FileSink> sink,
        std::shared_ptr<memory::MemoryPool> pool,
        uint64_t flushBufferSize)
        : sink_(std::move(sink)),
          pool_(std::move(pool)),
          flushBufferSize_(flushBufferSize),
          buf_(std::make_unique<dwio::common::DataBuffer<char> >(*pool_)) {
        reserveBuffer();
    }

    BufferedWriterSink::~BufferedWriterSink() {
        POLLUX_CHECK_EQ(
            buf_->size(),
            0,
            "Unexpected buffer data on BufferedWriterSink destruction");
    }

    void BufferedWriterSink::write(char value) {
        write(&value, 1);
    }

    void BufferedWriterSink::write(const char *data, uint64_t size) {
        // TODO Add logic for when size is larger than flushCount_
        POLLUX_CHECK_GE(
            flushBufferSize_,
            size,
            "write data size exceeds flush buffer size limit");

        if (buf_->size() + size > flushBufferSize_) {
            flush();
        }
        buf_->append(buf_->size(), data, size);
    }

    void BufferedWriterSink::flush() {
        if (buf_->size() == 0) {
            return;
        }

        sink_->write(std::move(*buf_));
        reserveBuffer();
    }

    void BufferedWriterSink::abort() {
        // TODO Add a flag to indicate sink is aborted to
        //  prevent flush and write operations after aborted.
        buf_->clear();
        sink_->close();
    }

    void BufferedWriterSink::close() {
        flush();
        buf_->clear();
        sink_->close();
    }

    void BufferedWriterSink::reserveBuffer() {
        POLLUX_CHECK_NOT_NULL(buf_);
        POLLUX_CHECK_EQ(buf_->size(), 0);
        POLLUX_CHECK_EQ(buf_->capacity(), 0);
        buf_->reserve(flushBufferSize_);
    }
} // namespace kumo::pollux::text
